import numpy as np
import os
import pickle
import matplotlib.pyplot as plt

model_select_index = {'cnn_nor': 5, 'cnn_unnor': 0, 'vi_nor': 12, 'vi_unnor': 12}

work_dir = os.getcwd()
pv_calculation_source_path = work_dir + '/result/detailed_soapy_test_result'
pv_calculation_result_path = work_dir + '/result/pv_calculation_for_detailed_soapy_test_result_supplyment'

cnn_nor_source = pv_calculation_source_path + '/CNN/normalized/5'
cnn_unnor_source = pv_calculation_source_path + '/CNN/unnormalized/0'
vi_nor_source = pv_calculation_source_path + '/Vision Transformer/normalized/12'
vi_unnor_source = pv_calculation_source_path + '/Vision Transformer/unnormalized/12'
model_source_path = [cnn_nor_source, cnn_unnor_source, vi_nor_source, vi_unnor_source]

cnn_nor_result = pv_calculation_result_path + '/CNN/normalized/5'
cnn_unnor_result = pv_calculation_result_path + '/CNN/unnormalized/0'
vi_nor_result = pv_calculation_result_path + '/Vision Transformer/normalized/12'
vi_unnor_result = pv_calculation_result_path + '/Vision Transformer/unnormalized/12'
model_result_path = [cnn_nor_result, cnn_unnor_result, vi_nor_result, vi_unnor_result]

r0_list = [str(round(0.05 + i * 0.01, 2)) for i in range(16)]
mag_list = [str(i) for i in range(8, 19)]


# define function for the processing of data to obtain the descriptive information of it
def calculate_descriptive(data):
    """"
    the output sequence of outcome is mean, std, media, 25_percentile, 75_percentile, max and min
    """
    data_mean = np.mean(data)
    data_std = np.std(data)
    data_media = np.median(data)
    data_25_percentile = np.percentile(data, 25)
    data_75_percentile = np.percentile(data, 75)
    data_max = np.max(data)
    data_min = np.min(data)
    return {'mean': data_mean, 'std': data_std, 'media': data_media, 'P_25': data_25_percentile,
            'P_75': data_75_percentile, 'max': data_max, 'min': data_min}


for i in range(4):
    # firstly calculate pv value under different r0
    pv_for_different_r0 = dict()
    for j in r0_list:
        work_source_path = model_source_path[i] + '/different_r0/r0=' + j

        pv_for_specific_r0 = dict()
        pv_for_specific_r0['pv_before_with_radius'] = []
        pv_for_specific_r0['pv_after_with_radius'] = []

        # make directory for specific r0
        work_result_path = model_result_path[i] + '/different_r0'
        if not os.path.exists(work_result_path + '/r0=' + j):
            os.makedirs(work_result_path + '/r0=' + j)

        for z in range(1000):
            phase_before = np.load(work_source_path + '/' + str(z) + '/phase_before_radius.npy', allow_pickle=True)
            phase_after = np.load(work_source_path + '/' + str(z) + '/phase_after_radius.npy', allow_pickle=True)

            pv_before_with_radius = np.max(phase_before) - np.min(phase_before)
            pv_after_with_radius = np.max(phase_after) - np.min(phase_after)

            # initiate variable for the storing of information under specific sample
            specific_sample_dirt = dict()
            specific_sample_dirt['pv_before_with_radius'] = pv_before_with_radius
            specific_sample_dirt['pv_after_with_radius'] = pv_after_with_radius

            # make directory for specific sample
            if not os.path.exists(work_result_path + '/r0=' + j + '/' + str(z)):
                os.makedirs(work_result_path + '/r0=' + j + '/' + str(z))

            # save information for specific sample
            with open(work_result_path + '/r0=' + j + '/' + str(z) + '/specific_sample_dirt', 'wb') as file:
                pickle.dump(specific_sample_dirt, file)

            # save information for specific sample in txt format
            with open(work_result_path + '/r0=' + j + '/' + str(z) + '/specific_sample_dirt.txt', 'w') as file:
                file.write(f'specific_sample_dirt = {specific_sample_dirt}')

            pv_for_specific_r0['pv_before_with_radius'].append(pv_before_with_radius)
            pv_for_specific_r0['pv_after_with_radius'].append(pv_after_with_radius)

            # print essential information on screen to tell the advancement of program
            print(f'model_index = {i}, r0 = {j}, element_index = {z}')

        # make directory for the saving of overall result under specific r0
        if not os.path.exists(model_result_path[i] + '/different_r0/r0=' + j + '/overall_result'):
            os.makedirs(model_result_path[i] + '/different_r0/r0=' + j + '/overall_result')

        specific_r0_overall_result_path = model_result_path[i] + '/different_r0/r0=' + j + '/overall_result'

        # save and plot information after processing
        with open(specific_r0_overall_result_path + '/overall_result_direc', 'wb') as file:
            pickle.dump(pv_for_specific_r0, file)
        plt.subplot(1, 2, 1)
        plt.hist(pv_for_specific_r0['pv_before_with_radius'])
        plt.xlabel('radius')
        plt.ylabel('count')
        plt.title('pv_before_with_radius')
        plt.subplot(1, 2, 2)
        plt.hist(pv_for_specific_r0['pv_after_with_radius'])
        plt.xlabel('radius')
        plt.ylabel('count')
        plt.title('pv_after_with_radius')
        plt.tight_layout()
        plt.savefig(specific_r0_overall_result_path + '/overall_result')
        plt.close('all')

        # store information under specific r0 to relevant directory
        pv_for_different_r0['r0=' + j] = pv_for_specific_r0

    # save overall information for all r0
    overall_result_path_r0 = model_result_path[i] + '/different_r0/overall_result'
    if not os.path.exists(model_result_path[i] + '/different_r0/overall_result'):
        os.makedirs(model_result_path[i] + '/different_r0/overall_result')
    with open(overall_result_path_r0 + '/overall', 'wb') as file:
        pickle.dump(pv_for_different_r0, file)
    with open(overall_result_path_r0 + '/overall.txt', 'w') as file:
        file.write(f'pv_for_different_r0 = {pv_for_different_r0}')

    # prepare and save descriptive information for all r0
    # Note!!!!! Instead of re-establish a completely new variable,
    # we just modify the variable storing the original information

    for j in pv_for_different_r0:

        pv_before_descriptive = calculate_descriptive(pv_for_different_r0[j]['pv_before_with_radius'])
        pv_for_different_r0[j]['pv_before_with_radius'] = pv_before_descriptive

        pv_after_descriptive = calculate_descriptive(pv_for_different_r0[j]['pv_after_with_radius'])
        pv_for_different_r0[j]['pv_after_with_radius'] = pv_after_descriptive

    with open(overall_result_path_r0 + '/descriptive_info', 'wb') as file:
        pickle.dump(pv_for_different_r0, file)
    with open(overall_result_path_r0 + '/descriptive_info.txt', 'w') as file:
        file.write(f'pv_for_different_r0 = {pv_for_different_r0}')

    # secondly calculate pv value under different magnitude
    pv_for_different_mag = dict()
    for j in mag_list:
        work_source_path = model_source_path[i] + '/different_mag/star_mag=' + j

        pv_for_specific_mag = dict()
        pv_for_specific_mag['pv_before_with_radius'] = []
        pv_for_specific_mag['pv_after_with_radius'] = []

        # make directory for specific magnitude
        work_result_path = model_result_path[i] + '/different_mag'
        if not os.path.exists(work_result_path + '/star_mag=' + j):
            os.makedirs(work_result_path + '/star_mag=' + j)

        for z in range(1000):
            phase_before = np.load(work_source_path + '/' + str(z) + '/phase_before_radius.npy', allow_pickle=True)
            phase_after = np.load(work_source_path + '/' + str(z) + '/phase_after_radius.npy', allow_pickle=True)

            pv_before_with_radius = np.max(phase_before) - np.min(phase_before)
            pv_after_with_radius = np.max(phase_after) - np.min(phase_after)

            # initiate variable for the storing of information under specific sample
            specific_sample_dirt = dict()
            specific_sample_dirt['pv_before_with_radius'] = pv_before_with_radius
            specific_sample_dirt['pv_after_with_radius'] = pv_after_with_radius

            # make directory for specific sample
            if not os.path.exists(work_result_path + '/star_mag=' + j + '/' + str(z)):
                os.makedirs(work_result_path + '/star_mag=' + j + '/' + str(z))

            # save information for specific sample
            with open(work_result_path + '/star_mag=' + j + '/' + str(z) + '/specific_sample_dirt', 'wb') as file:
                pickle.dump(specific_sample_dirt, file)

            # save information for specific sample in txt format
            with open(work_result_path + '/star_mag=' + j + '/' + str(z) + '/specific_sample_dirt.txt', 'w') as file:
                file.write(f'specific_sample_dirt = {specific_sample_dirt}')

            pv_for_specific_mag['pv_before_with_radius'].append(pv_before_with_radius)
            pv_for_specific_mag['pv_after_with_radius'].append(pv_after_with_radius)

            # print essential information on screen to tell the advancement of program
            print(f'model_index = {i}, magnitude = {j}, element_index = {z}')

        # make directory for the saving of overall result under specific magnitude
        if not os.path.exists(model_result_path[i] + '/different_mag/star_mag=' + j + '/overall_result'):
            os.makedirs(model_result_path[i] + '/different_mag/star_mag=' + j + '/overall_result')

        specific_mag_overall_result_path = model_result_path[i] + '/different_mag/star_mag=' + j + '/overall_result'

        # save and plot information after processing
        with open(specific_mag_overall_result_path + '/overall_result_direc', 'wb') as file:
            pickle.dump(pv_for_specific_mag, file)
        plt.subplot(1, 2, 1)
        plt.hist(pv_for_specific_mag['pv_before_with_radius'])
        plt.xlabel('radius')
        plt.ylabel('count')
        plt.title('pv_before_with_radius')
        plt.subplot(1, 2, 2)
        plt.hist(pv_for_specific_mag['pv_after_with_radius'])
        plt.xlabel('radius')
        plt.ylabel('count')
        plt.title('pv_after_with_radius')
        plt.tight_layout()
        plt.savefig(specific_mag_overall_result_path + '/overall_result')
        plt.close('all')

        # store information under specific magnitude to relevant directory
        pv_for_different_mag['star_mag=' + j] = pv_for_specific_mag

    # save overall information for all magnitude
    overall_result_path_mag = model_result_path[i] + '/different_mag/overall_result'
    if not os.path.exists(model_result_path[i] + '/different_mag/overall_result'):
        os.makedirs(model_result_path[i] + '/different_mag/overall_result')
    with open(overall_result_path_mag + '/overall', 'wb') as file:
        pickle.dump(pv_for_different_mag, file)
    with open(overall_result_path_mag + '/overall.txt', 'w') as file:
        file.write(f'pv_for_different_mag = {pv_for_different_mag}')

    # prepare and save descriptive information for all magnitude
    # Note!!!!! Instead of re-establish a completely new variable,
    # we just modify the variable storing the original information

    for j in pv_for_different_mag:
        pv_before_descriptive = calculate_descriptive(pv_for_different_mag[j]['pv_before_with_radius'])
        pv_for_different_mag[j]['pv_before_with_radius'] = pv_before_descriptive

        pv_after_descriptive = calculate_descriptive(pv_for_different_mag[j]['pv_after_with_radius'])
        pv_for_different_mag[j]['pv_after_with_radius'] = pv_after_descriptive

    with open(overall_result_path_mag + '/descriptive_info', 'wb') as file:
        pickle.dump(pv_for_different_mag, file)
    with open(overall_result_path_mag + '/descriptive_info.txt', 'w') as file:
        file.write(f'pv_for_different_mag = {pv_for_different_mag}')
